package LibDL.recommender.eval;

import LibDL.core.*;
import net.librec.recommender.item.KeyValue;

import java.util.List;
import java.util.ArrayList;
import java.util.Collections;
import java.lang.Math;

/*
* groundTruthList是测试值
* recList是自己模型得出的推荐结果
* topk是指需要评估推荐结果中前k个
*
* List<Tensor>是指每一个用户都有一个类型为Tensor的推荐结果Tensor里是item的id
* List<KeyValue<Tensor,Tensor>>中，每个用户一个KeyValue，KeyValue中第一个Tensor是item，整型；第二个是item对应的分值，double类型
* */

public class eval{
    public static boolean contain(Tensor tensor, Scalar tar){
        int eq = functional.eq(tensor, tar).nonzero().size(0);
        return eq != 0;
    }
    private static double ln2 = 0.693147181;
    public static double ndcg(List<KeyValue<Tensor, Tensor>> groundTruthList, List<Tensor> recList, int topk){
        double nDCG = 0.0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx).getKey();
            if(truth.size(0) > 0){
                StdVectorDouble val = groundTruthList.get(contextIdx).getValue().tolist_double();
                Tensor rec = recList.get(contextIdx);
                boolean hasdcgsValue = false;

                // calculate DCG
                double dcg = 0.0;
                ArrayList idcgsValue = new ArrayList();
                int topK = topk <= rec.size(0) ? topk : rec.size(0);

                for (int indexOfKey = 0; indexOfKey < topK; ++indexOfKey) {
                    Scalar itemID = rec.get(indexOfKey).item();

                    Tensor nonzero = functional.equal_indices(truth, itemID).nonzero();
                    if (nonzero.size(0) != 0) {
                        int index = (int)nonzero.item().to_long();

                        double rankvalue = val.get(index);
                        hasdcgsValue = true;
                        dcg += rankvalue / (Math.log(indexOfKey + 2) / ln2);//Maths.log(indexOfKey + 2, 2);

                        idcgsValue.add(rankvalue);
                    }
                }

                if(!hasdcgsValue||dcg == 0) {
                    ++nonZeroContext;
                    continue;
                }

                // calculate iDCG
                double idcg = 0.0d;
                Collections.sort(idcgsValue, Collections.reverseOrder());

                for(int i=0; i<idcgsValue.size(); i++) {
                    idcg += (double)idcgsValue.get(i) / (Math.log(i + 2) / ln2);//Maths.log(i + 2, 2);
                }

                if(idcg==0){
                    ++nonZeroContext;
                    continue;
                }
                nDCG += dcg / idcg;
                ++nonZeroContext;
            }
        }

        return nonZeroContext > 0 ? nDCG / nonZeroContext : 0.0d;
    }
    public static double hitrate(List<Tensor> groundTruthList, List<Tensor> recList, int topk){
        int totalHits = 0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            int truthSize = truth.size(0);
            if(truthSize == 0){
                continue;
            }
            truth = functional.topk(truth, 1).getFirst();
            if (truthSize == 1) {
                Tensor rec = recList.get(contextIdx);
                //topk
                rec = functional.topk(rec, topk).getFirst();
                if(contain(rec, truth.item())){
                    totalHits++;
                }
                nonZeroContext++;

            } else if (truth.size(0) > 1) {
                throw new IndexOutOfBoundsException("It is not a leave-one-out validation method! Please use leave-one-out validation method");
            }
        }

        return nonZeroContext > 0 ? 1.0 * totalHits / nonZeroContext : 0.0d;
    }
    public static double auc(List<Tensor> groundTruthList, List<Tensor> recList, int topk, int[] numDroppedArray){
        double auc = 0.0d;

        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        if (numDroppedArray == null || numDroppedArray.length != numContext){
            //throw new LibrecRuntimeException("please set rec.eval.auc.dropped.num arrays, length of numDroppedArray must be cardinality of groundTruthList.");
        }

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            if(truth.size(0) > 0){
                nonZeroContext++;
                Tensor rec = recList.get(contextIdx);
                int topK = rec.size(0);
                if(topk < topK){
                    topK = topk;
                    rec = functional.topk(rec, topK).getFirst();
                }
                int numDroppedItems = numDroppedArray[contextIdx] - topK;
                int numRelevantKeys = 0, numMissKeys = 0;
                for(int i = 0; i < topK; i++){
                    if(contain(truth, rec.get(i).item())){
                        numRelevantKeys++;
                    }else{
                        numMissKeys++;
                    }
                }
                //todo Drop
                int numEvaluatingItems =  numDroppedItems + topK;
                int numEvaluatingKeyValues = (numEvaluatingItems - numRelevantKeys) * numRelevantKeys;

                if (numEvaluatingKeyValues < 0) {
                    throw new IndexOutOfBoundsException("numEvaluatingKeyValues cannot be less than 0.");
                }

                //正样本数*负样本数为0
                if (numEvaluatingKeyValues == 0) {
                    auc += 0.5;
                    continue;
                }

                int numCorrectKeyValues = 0;
                int hits = 0;
                for (int i = 0; i < truth.size(0); i++) {
                    if (!contain(rec, truth.get(i).item())) {
                        numCorrectKeyValues += hits;
                    } else {
                        hits++;
                    }
                }

                numCorrectKeyValues += hits * (numDroppedItems - numMissKeys);

                auc += (numCorrectKeyValues + 0.0) / numEvaluatingKeyValues;
            }
        }

        return nonZeroContext > 0 ? auc / nonZeroContext : 0.0d;
    }
    public static double average_precision(List<Tensor> groundTruthList, List<Tensor> recList, int topk){
        double totalPrecision = 0.0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            if (truth.size(0) > 0) {
                Tensor rec = recList.get(contextIdx);

                int numHits = 0;
                int topK = topk < rec.size(0) ? topk : rec.size(0);
                double tempPrecision = 0.0d;
                for (int indexOfKey = 0; indexOfKey < topK; ++indexOfKey) {
                    Scalar itemID = rec.get(indexOfKey).item();
                    if (contain(truth, itemID)) {
                        numHits++;
                        tempPrecision += 1.0 * numHits / (indexOfKey + 1);
                    }
                }
                if(topK != 0) {
                    totalPrecision += tempPrecision / (truth.size(0) < topK ? truth.size(0) : topK);
                    nonZeroContext++;
                }
            }
        }
        return nonZeroContext > 0 ? totalPrecision / nonZeroContext : 0.0d;
    }
    public static double precision(List<Tensor> groundTruthList, List<Tensor> recList, int topk){
        double totalPrecision = 0.0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            if (truth.size(0) > 0) {
                Tensor rec = recList.get(contextIdx);

                int numHits = 0;
                int topK = topk < rec.size(0) ? topk : rec.size(0);
                for (int indexOfKey = 0; indexOfKey < topK; ++indexOfKey) {
                    Scalar itemID = rec.get(indexOfKey).item();
                    if (contain(truth, itemID)) {
                        numHits++;
                    }
                }
                totalPrecision += numHits / (topk + 0.0);
                nonZeroContext++;
            }
        }
        return nonZeroContext > 0 ? totalPrecision / nonZeroContext : 0.0d;
    }
    public static double reciprocal_rank(List<Tensor> groundTruthList, List<Tensor> recList, int topk){
        double reciprocalRank = 0.0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            if (truth.size(0) > 0) {
                Tensor rec = recList.get(contextIdx);

                int topK = topk < rec.size(0) ? topk : rec.size(0);
                for (int indexOfKey = 0; indexOfKey < topK; ++indexOfKey) {
                    Scalar itemID = rec.get(indexOfKey).item();
                    if (contain(truth, itemID)) {
                        reciprocalRank += 1.0 / (indexOfKey + 1);
                        break;
                    }
                }
                nonZeroContext++;
            }
        }
        return nonZeroContext > 0 ? reciprocalRank / nonZeroContext : 0.0d;
    }
    public static double average_reciprocal_hit_rank(List<Tensor> groundTruthList, List<Tensor> recList, int topk){
        double reciprocalRank = 0.0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            if (truth.size(0) > 0) {
                Tensor rec = recList.get(contextIdx);
                int topK = rec.size(0);
                if(topk < topK){
                    topK = topk;
                    rec = functional.topk(rec, topK).getFirst();
                }
                //todo
                Scalar itemID = truth.get(0).item();

                Tensor nonzero = functional.equal_indices(rec, itemID).nonzero();
                if (nonzero.size(0) != 0) {
                    int index = (int)nonzero.item().to_long();
                    reciprocalRank += 1.0 / (index + 1);
                }
                nonZeroContext++;
            }
        }
        return nonZeroContext > 0 ? reciprocalRank / nonZeroContext : 0.0d;
    }
    public static double recall(List<Tensor> groundTruthList, List<Tensor> recList, int topk){
        double totalRecall = 0.0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            if (truth.size(0) > 0) {
                Tensor rec = recList.get(contextIdx);

                int numHits = 0;
                int topK = topk < rec.size(0) ? topk : rec.size(0);
                for (int indexOfKey = 0; indexOfKey < topK; ++indexOfKey) {
                    Scalar itemID = rec.get(indexOfKey).item();
                    if (contain(truth, itemID)) {
                        numHits++;
                    }
                }
                totalRecall += numHits / (truth.size(0) + 0.0);
                nonZeroContext++;
            }
        }
        return nonZeroContext > 0 ? totalRecall / nonZeroContext : 0.0d;
    }
    public static double F1score(List<Tensor> groundTruthList, List<Tensor> recList, int topk){
        double totalF1Score = 0.0;
        int numContext = groundTruthList.size();
        int nonZeroContext = 0;

        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            if (truth.size(0) > 0) {
                Tensor rec = recList.get(contextIdx);

                int numHits = 0;
                int topK = topk < rec.size(0) ? topk : rec.size(0);
                for (int indexOfKey = 0; indexOfKey < topK; ++indexOfKey) {
                    Scalar itemID = rec.get(indexOfKey).item();
                    if (contain(truth, itemID)) {
                        numHits++;
                    }
                }
                double  recall = numHits / (truth.size(0) + 0.0);
                double precision =  numHits / (topk + 0.0);
                //todo
                totalF1Score += numHits != 0 ? (2 * (precision * recall) / (precision + recall)) : 0.0;
                nonZeroContext++;
            }
        }
        return nonZeroContext > 0 ? totalF1Score / nonZeroContext : 0.0d;
    }


    public static double mae(List<Tensor> groundTruthList, List<Tensor> recList){
        double mae = 0.0;
        int testSize = 0;
        int numContext = groundTruthList.size();
        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            Tensor rec = recList.get(contextIdx);
            if(rec.size(0) != truth.size(0)){
                throw new IndexOutOfBoundsException("num of recommendedList does not equal testMatrix num");
            }
            mae += truth.sub(rec).abs().sum(0).item().to_long();
            testSize += groundTruthList.get(contextIdx).size(0);
        }

        return testSize > 0 ? mae / testSize : 0.0d;
    }
    public static double rmse(List<Tensor> groundTruthList, List<Tensor> recList){
        double mae = 0.0;
        int testSize = 0;
        int numContext = groundTruthList.size();
        Scalar exp_2 = new Scalar(2);
        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            Tensor rec = recList.get(contextIdx);
            if(rec.size(0) != truth.size(0)){
                throw new IndexOutOfBoundsException("num of recommendedList does not equal testMatrix num");
            }
            mae += truth.sub(rec).pow(exp_2).sum(0).item().to_long();
            testSize += groundTruthList.get(contextIdx).size(0);
        }

        return testSize > 0 ? Math.sqrt(mae / testSize) : 0.0d;
    }
    public static double mse(List<Tensor> groundTruthList, List<Tensor> recList){
        double mae = 0.0;
        int testSize = 0;
        int numContext = groundTruthList.size();
        Scalar exp_2 = new Scalar(2);
        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            Tensor rec = recList.get(contextIdx);
            if(rec.size(0) != truth.size(0)){
                throw new IndexOutOfBoundsException("num of recommendedList does not equal testMatrix num");
            }
            mae += truth.sub(rec).pow(exp_2).sum(0).item().to_long();
            testSize += groundTruthList.get(contextIdx).size(0);
        }

        return testSize > 0 ? mae / testSize : 0.0d;
    }
    public static double mpe(List<Tensor> groundTruthList, List<Tensor> recList, double mpe){
        int mpeNum = 0;
        int testSize = 0;
        int numContext = groundTruthList.size();
        Scalar mpe_scalar = new Scalar(mpe);
        for (int contextIdx = 0; contextIdx < numContext; ++contextIdx) {
            Tensor truth = groundTruthList.get(contextIdx);
            Tensor rec = recList.get(contextIdx);
            if(rec.size(0) != truth.size(0)){
                throw new IndexOutOfBoundsException("num of recommendedList does not equal testMatrix num");
            }
            mpeNum += functional.gt(truth.sub(rec).abs(), mpe_scalar).sum(0).item().to_long();
            testSize += groundTruthList.get(contextIdx).size(0);
        }

        return testSize > 0 ? (mpeNum + 0.0) / testSize : 0.0d;
    }
}